#ifndef SRC_MPI_GET_MATRIX_BLK
#define SRC_MPI_GET_MATRIX_BLK

#include <NTL/mat_ZZ.h>

/**
 * @brief Get the submatrix `blk` from the matrix `mat` for process(i, j)
 *
 * @tparam T
 * @param mat           matrix `mat`: mat_row * mat_col
 * @param mat_row       # of mat's rows
 * @param mat_col       # of mat's columns
 * @param pi            index-i for the process grid
 * @param pj            index-j for the process grid
 * @param blk_row       # of blk's rows
 * @param blk_col       # of blk's columns
 * @return NTL::mat_ZZ
 */
template <typename T>
NTL::mat_ZZ get_blk(const T *mat, long mat_row, long mat_col, int pi, int pj,
                    long blk_row, long blk_col) {
  NTL::mat_ZZ blk;
  blk.SetDims(blk_row, blk_col);
  for (int ii = 0; ii < blk_row; ii++) {
    for (int jj = 0; jj < blk_col; jj++) {
      if (pi * blk_row + ii >= mat_row || pj * blk_col + jj >= mat_col) {
        continue;
      }
      blk[ii][jj] = mat[(pi * blk_row + ii) * mat_col + (pj * blk_col + jj)];
    }
  }
  return std::move(blk);
}

/**
 * @brief Get the submatrix `blk` from the submatrix of `mat` for process(i, j)
 *
 * @tparam T
 * @param mat           matrix `mat`: mat_row * mat_col
 * @param mat_row       # of mat's rows
 * @param mat_col       # of mat's columns
 * @param pi            index-i for the process grid
 * @param pj            index-j for the process grid
 * @param blk_row       # of blk's rows
 * @param blk_col       # of blk's columns
 * @param blki          index-i for submatrix grids of mat
 * @param blkj          index-j for submatrix grids of mat
 * @param bb_row
 * @param bb_col
 * @return NTL::mat_ZZ
 */
template <typename T>
NTL::mat_ZZ get_blk(const T *mat, long mat_row, long mat_col, int pi, int pj,
                    long blk_row, long blk_col, int blki, int blkj, long bb_row,
                    long bb_col) {
  NTL::mat_ZZ blk;
  blk.SetDims(bb_row, bb_col);
  for (int ii = 0; ii < bb_row; ii++) {
    for (int jj = 0; jj < bb_col; jj++) {
      const auto local_i = blki * bb_row + ii;
      const auto local_j = blkj * bb_col + jj;
      const auto global_i = pi * blk_row + local_i;
      const auto global_j = pj * blk_col + local_j;
      if (global_i >= mat_row || global_j >= mat_col) {
        continue;
      }
      if (local_i >= blk_row || local_j >= blk_col) {
        blk[ii][jj] = 0l;
      } else {
        blk[ii][jj] = mat[global_i * mat_col + global_j];
      }
    }
  }
  return std::move(blk);
}

#endif /* SRC_MPI_GET_MATRIX_BLK */
